package com.bw.realtime.common.function;

import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.streaming.api.windowing.triggers.Trigger;
import org.apache.flink.streaming.api.windowing.triggers.TriggerResult;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;

public class TimeAndCountTrigger extends Trigger<Object, TimeWindow> {
    private final long maxCount;
    private final long interval;

    public TimeAndCountTrigger(long maxCount, long interval) {
        this.maxCount = maxCount;
        this.interval = interval;
    }

    @Override
    public TriggerResult onElement(Object element, long timestamp, TimeWindow window, TriggerContext ctx) throws Exception {
        // 注册下一个周期性触发点
        long currentTime = ctx.getCurrentProcessingTime();
        long nextFireTimestamp = ((currentTime / interval) + 1) * interval;
        ctx.registerProcessingTimeTimer(nextFireTimestamp);
        
        // 更新元素计数状态
        ValueState<Long> count = ctx.getPartitionedState(
            new ValueStateDescriptor<>("count", Long.class, 0L));
        long newCount = count.value() + 1;
        count.update(newCount);
        
        // 当元素数量达到阈值时触发计算
        if (newCount >= maxCount) {
            count.clear();
            return TriggerResult.FIRE;
        }
        return TriggerResult.CONTINUE;
    }

    @Override
    public TriggerResult onProcessingTime(long time, TimeWindow window, TriggerContext ctx) throws Exception {
        // 周期性触发计算
        ValueState<Long> count = ctx.getPartitionedState(
            new ValueStateDescriptor<>("count", Long.class));
        count.clear();
        return TriggerResult.FIRE;
    }

    @Override
    public TriggerResult onEventTime(long time, TimeWindow window, TriggerContext ctx) throws Exception {
        return TriggerResult.CONTINUE;
    }

    @Override
    public void clear(TimeWindow window, TriggerContext ctx) throws Exception {
        // 清理所有状态和计时器
        ValueState<Long> count = ctx.getPartitionedState(
            new ValueStateDescriptor<>("count", Long.class));
        count.clear();
        
        long currentTime = ctx.getCurrentProcessingTime();
        long nextFireTimestamp = ((currentTime / interval) + 1) * interval;
        ctx.deleteProcessingTimeTimer(nextFireTimestamp);
    }
}